from pettingzoo.atari import (
    basketball_pong_v3,
    boxing_v2,
    combat_plane_v2,
    combat_tank_v2,
    double_dunk_v3,
    entombed_competitive_v3,
    entombed_cooperative_v3,
    flag_capture_v2,
    foozpong_v3,
    ice_hockey_v2,
    joust_v3,
    mario_bros_v3,
    maze_craze_v3,
    othello_v3,
    pong_v3,
    quadrapong_v4,
    space_invaders_v2,
    space_war_v2,
    surround_v2,
    tennis_v3,
    video_checkers_v4,
    volleyball_pong_v3,
    warlords_v3,
    wizard_of_wor_v3,
)

atari_environments = {
    "atari/basketball_pong_v3": basketball_pong_v3,
    "atari/boxing_v2": boxing_v2,
    "atari/combat_tank_v2": combat_tank_v2,
    "atari/combat_plane_v2": combat_plane_v2,
    "atari/double_dunk_v3": double_dunk_v3,
    "atari/entombed_competitive_v3": entombed_competitive_v3,
    "atari/entombed_cooperative_v3": entombed_cooperative_v3,
    "atari/flag_capture_v2": flag_capture_v2,
    "atari/foozpong_v3": foozpong_v3,
    "atari/joust_v3": joust_v3,
    "atari/ice_hockey_v2": ice_hockey_v2,
    "atari/maze_craze_v3": maze_craze_v3,
    "atari/mario_bros_v3": mario_bros_v3,
    "atari/othello_v3": othello_v3,
    "atari/pong_v3": pong_v3,
    "atari/quadrapong_v4": quadrapong_v4,
    "atari/space_invaders_v2": space_invaders_v2,
    "atari/space_war_v2": space_war_v2,
    "atari/surround_v2": surround_v2,
    "atari/tennis_v3": tennis_v3,
    "atari/video_checkers_v4": video_checkers_v4,
    "atari/volleyball_pong_v3": volleyball_pong_v3,
    "atari/wizard_of_wor_v3": wizard_of_wor_v3,
    "atari/warlords_v3": warlords_v3,
}
